"""服务识别与版本检测模块"""
import asyncio
import socket
import re
import ssl
from typing import Dict, Any, List
import subprocess
import json

class ServiceDetector:
    """服务类型和版本检测器"""
    
    def __init__(self):
        self.service_signatures = {
            21: {
                'name': 'FTP',
                'probes': [b''],
                'patterns': [b'220.*FTP', b'220.*FileZilla']
            },
            22: {
                'name': 'SSH',
                'probes': [b''],
                'patterns': [b'SSH-2.0-OpenSSH', b'SSH-1.99-OpenSSH']
            },
            23: {
                'name': 'Telnet',
                'probes': [b''],
                'patterns': [b'telnet', b'login:']
            },
            25: {
                'name': 'SMTP',
                'probes': [b'EHLO test\r\n'],
                'patterns': [b'220.*SMTP', b'220.*mail']
            },
            53: {
                'name': 'DNS',
                'probes': [b''],
                'patterns': [b'DNS']
            },
            80: {
                'name': 'HTTP',
                'probes': [b'GET / HTTP/1.0\r\n\r\n'],
                'patterns': [b'HTTP/1\\.[01]', b'Server:']
            },
            110: {
                'name': 'POP3',
                'probes': [b''],
                'patterns': [b'\\+OK.*POP3', b'\\+OK.*pop3']
            },
            143: {
                'name': 'IMAP',
                'probes': [b''],
                'patterns': [b'\\* OK.*IMAP', b'\\* OK.*imap']
            },
            443: {
                'name': 'HTTPS',
                'probes': [b'GET / HTTP/1.0\r\n\r\n'],
                'patterns': [b'HTTP/1\\.[01]', b'Server:']
            },
            993: {
                'name': 'IMAPS',
                'probes': [b''],
                'patterns': [b'SSL.*IMAP']
            },
            995: {
                'name': 'POP3S',
                'probes': [b''],
                'patterns': [b'SSL.*POP3']
            },
            3306: {
                'name': 'MySQL',
                'probes': [b''],
                'patterns': [b'mysql', b'MySQL']
            },
            5432: {
                'name': 'PostgreSQL',
                'probes': [b''],
                'patterns': [b'postgresql', b'PostgreSQL']
            },
            6379: {
                'name': 'Redis',
                'probes': [b'PING\r\n'],
                'patterns': [b'\\+PONG']
            },
            27017: {
                'name': 'MongoDB',
                'probes': [b''],
                'patterns': [b'mongodb', b'MongoDB']
            }
        }
    
    async def detect_service(self, target: str, port: int, timeout: int) -> Dict[str, Any]:
        """检测服务类型和版本"""
        service_info = {
            'service': 'unknown',
            'version': 'unknown',
            'banner': '',
            'cpe': '',
            'confidence': 0
        }
        
        # 首先尝试基于端口号的快速识别
        if port in self.service_signatures:
            sig = self.service_signatures[port]
            service_info['service'] = sig['name']
            service_info['confidence'] = 30
        
        # 尝试获取服务banner
        banner = await self._get_banner(target, port, timeout)
        if banner:
            service_info['banner'] = banner.decode('utf-8', errors='ignore')
            
            # 基于banner进行服务识别
            detected = self._identify_from_banner(banner, port)
            if detected:
                service_info.update(detected)
        
        # 对于HTTP/HTTPS，进行更详细的检测
        if port in [80, 443, 8080, 8443]:
            http_info = await self._detect_http_service(target, port, timeout)
            if http_info:
                service_info.update(http_info)
        
        # 对于SSL/TLS服务，获取证书信息
        if port in [443, 993, 995, 8443]:
            ssl_info = await self._detect_ssl_service(target, port, timeout)
            if ssl_info:
                service_info.update(ssl_info)
        
        return service_info
    
    async def _get_banner(self, target: str, port: int, timeout: int) -> bytes:
        """获取服务banner"""
        try:
            reader, writer = await asyncio.wait_for(
                asyncio.open_connection(target, port),
                timeout=timeout
            )
            
            # 发送探测数据
            if port in self.service_signatures:
                probes = self.service_signatures[port]['probes']
                for probe in probes:
                    if probe:
                        writer.write(probe)
                        await writer.drain()
            
            # 接收响应
            banner = await asyncio.wait_for(
                reader.read(1024),
                timeout=timeout
            )
            
            writer.close()
            await writer.wait_closed()
            
            return banner
            
        except Exception:
            return b''
    
    def _identify_from_banner(self, banner: bytes, port: int) -> Dict[str, Any]:
        """从banner中识别服务"""
        banner_str = banner.decode('utf-8', errors='ignore')
        
        # 服务识别模式
        patterns = {
            'SSH': {
                'pattern': r'SSH-(\d+\.\d+)-([^\r\n]+)',
                'service': 'SSH',
                'version_group': 2
            },
            'FTP': {
                'pattern': r'220[^\r\n]*FTP[^\r\n]*([^\r\n]+)',
                'service': 'FTP',
                'version_group': 1
            },
            'HTTP': {
                'pattern': r'Server: ([^\r\n]+)',
                'service': 'HTTP',
                'version_group': 1
            },
            'SMTP': {
                'pattern': r'220[^\r\n]*SMTP[^\r\n]*([^\r\n]+)',
                'service': 'SMTP',
                'version_group': 1
            },
            'MySQL': {
                'pattern': r'mysql[^\r\n]*([^\r\n]+)',
                'service': 'MySQL',
                'version_group': 1
            },
            'PostgreSQL': {
                'pattern': r'PostgreSQL[^\r\n]*([^\r\n]+)',
                'service': 'PostgreSQL',
                'version_group': 1
            }
        }
        
        for service_name, pattern_info in patterns.items():
            match = re.search(pattern_info['pattern'], banner_str, re.IGNORECASE)
            if match:
                version = match.group(pattern_info['version_group']).strip()
                return {
                    'service': pattern_info['service'],
                    'version': version,
                    'confidence': 80
                }
        
        return None
    
    async def _detect_http_service(self, target: str, port: int, 
                                 timeout: int) -> Dict[str, Any]:
        """检测HTTP服务详细信息"""
        try:
            # 构建HTTP请求
            request = (
                f"GET / HTTP/1.1\r\n"
                f"Host: {target}\r\n"
                f"User-Agent: LightPort/1.0\r\n"
                f"Accept: */*\r\n"
                f"\r\n"
            ).encode()
            
            reader, writer = await asyncio.wait_for(
                asyncio.open_connection(target, port),
                timeout=timeout
            )
            
            writer.write(request)
            await writer.drain()
            
            response = await asyncio.wait_for(
                reader.read(4096),
                timeout=timeout
            )
            
            writer.close()
            await writer.wait_closed()
            
            response_str = response.decode('utf-8', errors='ignore')
            
            # 解析HTTP响应
            info = {
                'service': 'HTTP',
                'confidence': 90
            }
            
            # 提取服务器信息
            server_match = re.search(r'Server: ([^\r\n]+)', response_str, re.IGNORECASE)
            if server_match:
                info['version'] = server_match.group(1).strip()
            
            # 提取其他HTTP头信息
            headers = {}
            for line in response_str.split('\r\n')[1:]:
                if ':' in line:
                    key, value = line.split(':', 1)
                    headers[key.strip()] = value.strip()
            
            info['http_headers'] = headers
            
            return info
            
        except Exception:
            return None
    
    async def _detect_ssl_service(self, target: str, port: int, 
                                timeout: int) -> Dict[str, Any]:
        """检测SSL/TLS服务信息"""
        try:
            context = ssl.create_default_context()
            context.check_hostname = False
            context.verify_mode = ssl.CERT_NONE
            
            reader, writer = await asyncio.wait_for(
                asyncio.open_connection(target, port, ssl=context),
                timeout=timeout
            )
            
            # 获取SSL证书信息
            cert = writer.get_extra_info('peercert')
            
            writer.close()
            await writer.wait_closed()
            
            if cert:
                return {
                    'ssl_cert': {
                        'subject': dict(x[0] for x in cert['subject']),
                        'issuer': dict(x[0] for x in cert['issuer']),
                        'version': cert['version'],
                        'serial_number': cert['serialNumber'],
                        'not_before': cert['notBefore'],
                        'not_after': cert['notAfter']
                    }
                }
            
        except Exception:
            pass
        
        return None
    
    def _generate_cpe(self, service: str, version: str) -> str:
        """生成CPE（通用平台枚举）标识符"""
        service_lower = service.lower()
        
        cpe_map = {
            'ssh': 'cpe:/a:openbsd:openssh',
            'http': 'cpe:/a:apache:http_server',
            'ftp': 'cpe:/a:filezilla:filezilla_server',
            'mysql': 'cpe:/a:oracle:mysql',
            'postgresql': 'cpe:/a:postgresql:postgresql'
        }
        
        base_cpe = cpe_map.get(service_lower, f'cpe:/a:{service_lower}:{service_lower}')
        
        if version and version != 'unknown':
            return f"{base_cpe}:{version}"
        
        return base_cpe
